package main

import (
	"fmt"
	"log"
	"sync"
	"time"
)

type ETLProcessor struct {
	extractWorkers   int
	transformWorkers int
	loadWorkers      int

	extractChan   chan *RawData
	transformChan chan *TransformedData
	loadChan      chan *TransformedData

	wg sync.WaitGroup
}

type RawData struct {
	ID   int
	Data string
}

type TransformedData struct {
	ID        int
	Data      string
	Processed string
	Timestamp time.Time
}

func NewETLProcessor(extract, transform, load int) *ETLProcessor {
	return &ETLProcessor{
		extractWorkers:   extract,
		transformWorkers: transform,
		loadWorkers:      load,
		extractChan:      make(chan *RawData, 1000),
		transformChan:    make(chan *TransformedData, 1000),
		loadChan:         make(chan *TransformedData, 1000),
	}
}

func (e *ETLProcessor) Start() {
	// 启动各个阶段的worker
	for i := 0; i < e.extractWorkers; i++ {
		go e.extractWorker(i)
	}

	for i := 0; i < e.transformWorkers; i++ {
		go e.transformWorker(i)
	}

	for i := 0; i < e.loadWorkers; i++ {
		go e.loadWorker(i)
	}
}

func (e *ETLProcessor) Process(dataSource []*RawData) {
	e.wg.Add(len(dataSource))

	// 提取阶段
	go func() {
		for _, data := range dataSource {
			e.extractChan <- data
		}
		close(e.extractChan)
	}()

	e.wg.Wait()
	close(e.transformChan)
	close(e.loadChan)
}

func (e *ETLProcessor) extractWorker(workerID int) {
	for rawData := range e.extractChan {
		// 模拟数据提取
		time.Sleep(5 * time.Millisecond)

		transformed := &TransformedData{
			ID:        rawData.ID,
			Data:      rawData.Data,
			Timestamp: time.Now(),
		}

		e.transformChan <- transformed
	}
}

func (e *ETLProcessor) transformWorker(workerID int) {
	for data := range e.transformChan {
		// 模拟数据转换
		time.Sleep(10 * time.Millisecond)

		data.Processed = fmt.Sprintf("TRANSFORMED-%s", data.Data)
		e.loadChan <- data
	}
}

func (e *ETLProcessor) loadWorker(workerID int) {
	for data := range e.loadChan {
		// 模拟数据加载
		time.Sleep(5 * time.Millisecond)

		fmt.Printf("Worker%d加载数据: ID=%d, Data=%s\n", workerID, data.ID, data.Processed)
		e.wg.Done()
	}
}

func main() {
	// 生成测试数据
	var testData []*RawData
	for i := 0; i < 10000; i++ {
		testData = append(testData, &RawData{
			ID:   i,
			Data: fmt.Sprintf("raw_data_%d", i),
		})
	}

	processor := NewETLProcessor(10, 20, 5) // 10个提取，20个转换，5个加载worker
	processor.Start()

	start := time.Now()
	processor.Process(testData)

	elapsed := time.Since(start)
	log.Printf("处理%d条数据耗时: %v", len(testData), elapsed)
}
